import signal
import sys
from huggingface_hub import HfApi
from typing import List, Tuple, Dict
from sys import exit
import os
import argparse
import asyncio
from downloader import DownloadManager
import json
import requests
from requests.auth import HTTPProxyAuth

class HF:
    def __init__(self, author: str, model_name: str, dataset_name:str, 
                 model_revision: str, dataset_revision: str,
                 token: str = None, proxies=None, base_url: str = "https://huggingface.co"):
        """
        初始化HF类
        Args:
            model_name: 模型名称
            dataset_name: 数据集名称
            model_revision: 模型版本
            dataset_revision: 数据集版本
            token: Hugging Face 访问令牌
            proxies: 代理配置，字典
            base_url: 基础URL，默认为 https://huggingface.co，可替换为其他镜像站如 https://hf-mirror.com
        """
        self.author = author
        self.model_name = model_name
        self.dataset_name = dataset_name
        self.model_revision = model_revision
        self.dataset_revision = dataset_revision
        self.token = token
        self.base_url = base_url        
        # 创建一个session来处理所有请求
        self.session = requests.Session()
        if proxies:
            self.session.proxies = proxies
        if token:
            self.session.headers.update({
                "Authorization": f"Bearer {token}",
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
            })

    def get_author_models(self, author: str) -> List[str]:
        """
        获取指定作者的所有模型列表
        """
        try:
            BASE_URL = f"{self.base_url}/api/models"
            url = f"{BASE_URL}?author={author}"
            response = self.session.get(url)
            response.raise_for_status()
            data = response.json()
            return [model['id'].split('/')[-1] for model in data]            
        except Exception as e:
            print(f"Error getting models for author {author}: {str(e)}")
            return []

    def get_author_datasets(self, author: str) -> List[str]:
        """
        获取指定作者的所有Dataset列表
        """
        try:
            BASE_URL = f"{self.base_url}/api/datasets"
            url = f"{BASE_URL}?author={author}"
            response = self.session.get(url)
            response.raise_for_status()
            data = response.json()
            return [dataset['id'].split('/')[-1] for dataset in data]            
        except Exception as e:
            print(f"Error getting datasets for author {author}: {str(e)}")
            return []
                
    def get_model_files(self) -> List[Tuple[str, str, str]]:
        """
        递归获取model的所有文件列表。
        :return: 文件信息列表，每个元素为 (url, <author>/<model_name>/dir, file_name)
        """
        return self._fetch_model_files(dir_path="")

    def get_dataset_files(self) -> List[Tuple[str, str, str]]:
        """
        递归获取dataset的所有文件列表。
        :return: 文件信息列表，每个元素为 (url, <author>/<dataset_name>/dir, file_name)
        """
        return self._fetch_dataset_files(dir_path="")
    
    def _fetch_model_files(self, dir_path : str ="") -> List[Tuple[str, str, str]]:
        """
        递归获取指定目录下的model所有文件信息。
        :param dir_path: 当前目录路径，默认为空表示根目录
        :return: 文件信息列表，每个元素为 (url, <author>/<model_name>/dir, file_name)
        """
        if dir_path:
            url = f"{self.base_url}/api/models/{self.author}/{self.model_name}/tree/{self.model_revision}/{dir_path}"
        else:
            url = f"{self.base_url}/api/models/{self.author}/{self.model_name}/tree/{self.model_revision}"
        response = self.session.get(url)
        if response.status_code != 200:
            print(f"Failed to fetch data from {url}. Status code: {response.status_code}")
            return []
        data = response.json()
        result = []
        for file_info in data:
            file_size = file_info['size']
            file_path = file_info['path']
            file_name = os.path.basename(file_path) # 文件名
            if file_info["type"] == "directory":  # 如果是目录，递归获取
                result.extend(self._fetch_model_files(file_path))
            else:  # 如果是文件，生成元组并加入结果列表
                full_url = f"{self.base_url}/{self.author}/{self.model_name}/resolve/{self.model_revision}/{file_path}"
                if dir_path:
                    relative_path = f"{self.author}/{self.model_name}/" + os.path.dirname(file_path)    
                else:
                    relative_path = f"{self.author}/{self.model_name}"
                result.append((full_url, relative_path, file_name, file_size))

        return result


    def _fetch_dataset_files(self, dir_path : str ="") -> List[Tuple[str, str, str]]:
        """
        递归获取指定目录下的dataset所有文件信息。
        :param dir_path: 当前目录路径，默认为空表示根目录
        :return: 文件信息列表，每个元素为 (url, <author>/<model_name>/dir, file_name)
        """
        if dir_path:
            url = f"{self.base_url}/api/datasets/{self.author}/{self.dataset_name}/tree/{self.dataset_revision}/{dir_path}"
        else:
            url = f"{self.base_url}/api/datasets/{self.author}/{self.dataset_name}/tree/{self.dataset_revision}"
        response = self.session.get(url)
        if response.status_code != 200:
            print(f"Failed to fetch data from {url}. Status code: {response.status_code}")
            return []
        data = response.json()
        result = []
        for file_info in data:
            file_size = file_info['size']
            file_path = file_info['path']
            file_name = os.path.basename(file_path) # 文件名
            if file_info["type"] == "directory":  # 如果是目录，递归获取
                result.extend(self._fetch_dataset_files(file_path))
            else:  # 如果是文件，生成元组并加入结果列表
                full_url = f"{self.base_url}/datasets/{self.author}/{self.dataset_name}/resolve/{self.dataset_revision}/{file_path}?download=true"
                if dir_path:
                    relative_path = f"{self.author}/{self.dataset_name}/" + os.path.dirname(file_path)    
                else:
                    relative_path = f"{self.author}/{self.dataset_name}"
                result.append((full_url, relative_path, file_name, file_size))

        return result  

class Cmd:
    def __init__(self):
        """
        初始化命令行参数解析器
        """
        self.parser = argparse.ArgumentParser(
            description='Hugging Face 工具',
            formatter_class=argparse.RawTextHelpFormatter
        )
        self.parser.add_argument(
            '-f', '--config', 
            type=str, 
            help='''配置文件路径。配置文件格式(JSON):
{
    "hf_token": "你的HuggingFace令牌",
    "base_url": "https://huggingface.co",
    "author": "作者名称",
    "model": "模型名称",
    "dataset": "数据集名称"
}
注意:
- hf_token 和 base_url 为必填项
- author 为必填项
- model 和 dataset 至少需要填写一个'''
        )
        self.config = None
        self.manager = None  #  manager 属性,用于停止下载任务

    def _print_help(self):
        """
        打印帮助信息
        """
        self.parser.print_help()
        print("\n使用示例:")
        print("1. 使用配置文件:")
        print("    hf -f hf_config.json")
        print("\n配置文件示例(config.json):")
        print('''{
    "hf_token": "你的HuggingFace令牌",
    "base_url": "https://huggingface.co",
    "author": "open-thoughts",
    "model": "",
    "dataset": "OpenThoughts-114k"
}''')

    def _load_config(self, config_file: str) -> Dict:
        """
        从配置文件加载参数        
        Args:
            config_file: 配置文件路径            
        Returns:
            Dict: 包含配置参数的字典
        """
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
                
                # 验证 hf_token 字段
                hf_token = config.get('hf_token')
                if not hf_token:
                    raise ValueError("配置文件中 'hf_token' 字段不能为空")
                    
                # 验证 base_url 字段
                base_url = config.get('base_url')
                if not base_url:
                    raise ValueError("配置文件中 'base_url' 字段不能为空")
                    
                # 验证 author 字段
                author = config.get('author')
                if not author:
                    raise ValueError("配置文件中 'author' 字段不能为空")
                
                # 验证 model 和 dataset 字段
                model = config.get('model')
                dataset = config.get('dataset')
                if not model and not dataset:
                    raise ValueError("配置文件中 'model' 和 'dataset' 字段不能同时为空")
                    
                # 设置下载配置的默认值
                if 'download' not in config:
                    config['download'] = {}
                if 'max_concurrent_files' not in config['download']:
                    config['download']['max_concurrent_files'] = 5
                if 'max_concurrent_chunks' not in config['download']:
                    config['download']['max_concurrent_chunks'] = 3
                if 'chunk_size' not in config['download']:
                    config['download']['chunk_size'] = 50*1024*1024  # 50MB
                
                # 处理代理配置
                proxy_config = config['download'].get('proxy', {})
                # 默认不使用代理
                use_proxy = proxy_config.get('use_proxy', False)
                
                if use_proxy:
                    from downloader import ProxyConfig
                    proxy_settings = proxy_config.get('settings')
                    if not proxy_settings:
                        raise ValueError("启用代理时，必须提供 'settings' 配置")
                    
                    if isinstance(proxy_settings, str):
                        # 如果是字符串，直接作为URL
                        config['download']['proxy'] = ProxyConfig(url=proxy_settings)
                    elif isinstance(proxy_settings, dict):
                        # 如果是字典，可能包含更多配置
                        proxy_url = proxy_settings.get('http') or proxy_settings.get('https')
                        if not proxy_url:
                            raise ValueError("代理配置必须包含 'http' 或 'https' URL")
                        config['download']['proxy'] = ProxyConfig(
                            url=proxy_url,
                            auth=proxy_settings.get('auth'),
                            headers=proxy_settings.get('headers')
                        )
                else:
                    config['download']['proxy'] = None
                    
                return config
                
        except FileNotFoundError:
            print(f"错误: 找不到配置文件 {config_file}")
            return {}
        except json.JSONDecodeError:
            print(f"错误: 配置文件 {config_file} 格式不正确")
            return {}
        except ValueError as e:
            print(f"错误: {str(e)}")
            return {}
        except Exception as e:
            print(f"错误: 读取配置文件时发生错误: {str(e)}")
            return {}

    def check_options(self) -> bool:
        """
        检查命令行参数和配置文件
        Returns:
            bool: 检查是否通过
        """
        args = self.parser.parse_args()
        
        # 如果没有任何参数，显示帮助信息
        if not args.config:
            self._print_help()
            return False

        try:
            # 检查配置文件是否存在
            if not os.path.exists(args.config):
                print(f"错误: 找不到配置文件 {args.config}")
                self._print_help()
                return False

            # 如果指定了配置文件，从配置文件读取参数
            config = self._load_config(args.config)
            if not config:
                self._print_help()
                return False
            
            self.config = config
            return True
            
        except Exception as e:
            print(f"错误: {str(e)}")
            self._print_help()
            return False

    async def process(self, hf: HF, manager: DownloadManager) -> None:
        """
        处理下载任务
        Args:
            hf: HF类实例
            manager: DownloadManager类实例
        """
        self.manager = manager  # 保存 manager 实例
        try:
            author = self.config.get('author')
            model = self.config.get('model')
            dataset = self.config.get('dataset')
            
            if model:
                print(f"\n开始下载 HuggingFace Model {author}/{model} 的文件 ...")
                files = hf.get_model_files()
                for url, directory, filename, total_size in files:
                    await manager.add_task(url, directory, filename, total_size)
                await manager.start_all()                
            elif dataset:
                print(f"\n开始下载 HuggingFace Dataset {author}/{dataset} 的文件 ...")
                files = hf.get_dataset_files()
                for url, directory, filename, total_size in files:
                    await manager.add_task(url, directory, filename, total_size)
                await manager.start_all()
            else:
                self._print_help()
                return
            
        except Exception as e:
            print(f"错误: {str(e)}")
            self._print_help()
            return

    async def handle_signal(self):
        """处理信号"""
        if self.manager:
            await self.manager.cancel_all_tasks()

    def get_config(self):
        return self.config
    

async def main():
    """主处理程序"""
    # 命令行处理类
    cmd = Cmd()
    
    # 检查命令行参数
    if not cmd.check_options():
        exit(1)
    
    # 获取配置
    config = cmd.get_config()
    
    # 初始化实例
    hf = HF(author=config['author'], 
            model_name=config['model'], 
            dataset_name=config['dataset'], 
            model_revision=config['model_revision'], 
            dataset_revision=config['dataset_revision'],
            token=config['hf_token'],
            proxies=None, 
            base_url=config['base_url'])

    # 配置下载管理器
    manager = DownloadManager(
        max_concurrent_files=config['download']['max_concurrent_files'],
        max_concurrent_chunks=config['download']['max_concurrent_chunks'],
        chunk_size=config['download']['chunk_size'],
        default_proxy=None,
        hf_token=config['hf_token']
    )

    # 添加基本的信号处理
    def signal_handler(sig, frame):
        print('\n收到中断信号，正在退出程序...')
        # 创建一个任务来处理取消操作
        loop = asyncio.get_running_loop()
        loop.create_task(cmd.handle_signal())
    
    # 设置信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    
    try:
        # 执行下载任务
        await cmd.process(hf, manager)
    finally:
        # 恢复默认信号处理
        signal.signal(signal.SIGINT, signal.default_int_handler)


# 主函数
if __name__ == "__main__":
    # Windows平台需要使用不同的事件循环策略
    if os.name == 'nt':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    
    try:
        # 运行主程序
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n程序已退出")
    except SystemExit:
        pass
    

